# -*- coding: utf-8 -*-
import matplotlib
matplotlib.use('TkAgg')  # 防止Qt/Wayland问题，指定后端
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from collections import deque
import pickle
import numpy as np
import random
import math

with open("data/data_floors_0427.pkl", "rb") as haddle:
    data_floors = pickle.load(haddle)

num_floors, \
floor_width, \
floor_length, \
floor_height, \
exits_on_ground, \
cell_size, \
speed_floor_down_radio, \
time_update, \
floors, \
stairs_up, \
stairs_down, \
exits, \
people = data_floors

# -----------------------------
# 4. 势能场计算（每层）
# -----------------------------
static_dist = np.full((num_floors, floor_width, floor_length), np.inf)

dq = deque()
for z,(i,j) in exits:
    dq.append((z,i,j))
    static_dist[z,i,j] = 0

while dq:
    cz,ci,cj = dq.popleft()
    #楼梯移动
    if floors[cz,ci,cj] == 4:
        dz = 1
        di = 0
        dj = 0
        nz, ni, nj = cz+dz, ci+di, cj+dj
        if 0<=ni<floor_width and 0<=nj<floor_length and 0<=nz<num_floors:            
            move_length = floor_height
            if floors[nz,ni,nj] in (0,3,4,5) and static_dist[nz,ni,nj]>static_dist[cz,ci,cj]+move_length:
                static_dist[nz,ni,nj] = static_dist[cz,ci,cj]+move_length
                dq.append((nz,ni,nj))
    #平层移动
    dz = 0
    for di,dj in [(-1,0),(1,0),(0,-1),(0,1),(-1,-1),(-1,1),(1,-1),(1,1)]:
        nz, ni, nj = cz+dz, ci+di, cj+dj
        if 0<=ni<floor_width and 0<=nj<floor_length and 0<=nz<num_floors:
            move_length = ((di)**2 + (dj)**2) ** 0.5
            if floors[nz,ni,nj] in (0,3,4,5) and static_dist[nz,ni,nj]>static_dist[cz,ci,cj]+move_length:
                static_dist[nz,ni,nj] = static_dist[cz,ci,cj]+move_length
                dq.append((nz,ni,nj))

# -----------------------------
# 5. 动态模拟部分
# -----------------------------
time_elapsed = 0  # 疏散时间计数器（单位秒）

def density_at(f, pos, people):
    # 计算当前格子附近密度
    count = sum(1 for p in people if p[0]==f and np.linalg.norm(np.array(p[1])-np.array(pos))<=1)
    area = np.pi * (1.5*cell_size)**2
    density = count / area
    return density

def get_velocity(density):
    if density >= 4:
        velocity = 0.033 * density**2 - 0.636* density + 3.362
    else:
        velocity = 1.4    
    return velocity

import time
start = time.time()
iter = 0
while len(people) > 0:
    iter += 1
    print ("run iter %d, len(people): %d" %(iter, len(people)))
    time_elapsed += time_update  # 每帧time_update秒 (每time_update秒更新一次)
    
    # 人员移动
    new_people = []
    floor_count = [0]*num_floors
    for p in people:
        f,(i,j),move_length = p

        if floors[f,i,j]==3 and f==0:
            continue  # 到达出口

        density = density_at(f, (i,j), people)
        velocity = get_velocity(density)

        # 楼梯移动过程，(f,i,j)是目标位置，还未到达，floors[f,i,j]是位置属性，目标位置的位置属性是向上的楼梯口
        if move_length>0 and move_length != math.inf and floors[f,i,j]==4:
            new_people.append([f,(i,j),move_length-velocity*speed_floor_down_radio*time_update])
            floor_count[f+1] +=1   #人还在上一层，所以统计上一层
            continue   

        # 平面移动过程，(f,i,j)是目标位置，还未到达，floors[f,i,j]是位置属性，目标位置的位置属性不是向上的楼梯口
        if move_length>0 and move_length != math.inf and floors[f,i,j]!=4:
            new_people.append([f,(i,j),move_length-velocity*time_update])
            floor_count[f] +=1     
            continue     

        # 楼梯移动距离计算，当前位置floors[f,i,j]的位置属性是向下的楼梯口
        if floors[f,i,j]==5:
            best_move = (i, j)
            move_length = floor_height
            new_people.append([f-1,(i,j),move_length])     #目标是下一层
            floor_count[f] +=1      #人还在当前层，所以统计当前层
            continue

        # 平面移动目标位置选择和距离计算
        min_dist = None
        best_move = (i,j)
        for di,dj in [(-1,0),(1,0),(0,-1),(0,1),(-1,-1),(-1,1),(1,-1),(1,1)]:
            ni,nj = i+di,j+dj
            if 0<=ni<floor_width and 0<=nj<floor_length:
                if floors[f,ni,nj] in (0,3,4,5):
                    if min_dist is None or static_dist[f,ni,nj]<min_dist:
                        min_dist = static_dist[f,ni,nj]
                        best_move = (ni,nj)
                    elif static_dist[f,ni,nj]==min_dist and random.random() <= 0.5:
                        min_dist = static_dist[f,ni,nj]
                        best_move = (ni,nj)                        
        ni,nj = best_move        
        move_length = ((ni-i)**2 + (nj-j)**2) ** 0.5
        new_people.append([f,(ni,nj),move_length])
        #当前层统计和绘图
        floor_count[f] +=1  
    
    people = new_people
        
print (f'(Static) Evacuation finished in {time_elapsed:.1f} seconds.')
end = time.time()
print ("计算时间", end - start)